package com.shxs.LLMConnect.LLMTools;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.shxs.LLMConnect.BlueLM.BlueLM;
import com.shxs.LLMConnect.LLMConnector;
import com.shxs.LLMConnect.chatGLM.ChatGLMConnector;
import com.shxs.model.entity.Dialog;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.*;

public class LLMSplitTool {

    LLMConnector connector;

    public LLMSplitTool(LLMConnector connector) {
        this.connector = connector;
        this.addVoiceActor("精英男生");
        this.addVoiceActor("霸道青年音色");
        this.addVoiceActor("青年大学生音色");
        this.addVoiceActor("男性主持人");
        this.addVoiceActor("男性有声书1");
        this.addVoiceActor("男性有声书2");
        this.addVoiceActor("少女音色");
        this.addVoiceActor("御姐音色");
        this.addVoiceActor("成熟女性音色");
        this.addVoiceActor("甜美女性音色");
        this.addVoiceActor("女性主持人");
        this.addVoiceActor("女性有声书1");
        this.addVoiceActor("女性有声书2");
    }

    List<String> voiceActor = new ArrayList<String>();

    public void addVoiceActor(String line) {
        voiceActor.add(line);
    }

    public int getVoiceActor(String line) {
        return voiceActor.indexOf(line);
    }
    String prompt = "你是一个配音角色分配机，我会给你一个台词本，其中所有带双引号的内容前，有符号[i]，i表示第几个符号，从0开始。通常符[i]表示接下来的内容为角色配音." +
            "你需要告诉我i对应的角色名，以格式“i：角色名：配音演员名”输出每行内容，必须严格遵照格式进行输出，不要输出任何其他内容，需要注意的是，对于第一人称小说不知道主角名字的情况下可以用“我”来表示主角名，如果是没有出现角色名的旁人角色，就用旁人男护着旁人女表示角色名，并且存在[i]后跟的不是角色台词而是特殊含义的词，此时就用“旁白”来表示，明白吗\n" +
            "\n" +
            "\n" +
            "请合理分配配音演员。注意由于配音演员数量有限，会有不同角色有相同配音演员的情况存在。";

    public void setPrompt(String prompt) {
        this.prompt = prompt;
    }
    public String doSplitWithHistory(String content,Map<String,String> historyRoleVoiceMap) {
        StringBuilder sb = new StringBuilder();
        sb.append(prompt);
        if (connector.getHistoryLength() == 0) {
            sb.append("配音演员有：" + voiceActor.toString() + "\n");
        }
        if(historyRoleVoiceMap!=null){
            System.out.println(historyRoleVoiceMap.toString());
            sb.append("当前角色和配音演员分配情况是：" + historyRoleVoiceMap.toString() + "\n");
        }
        sb.append("接下来是正文：");
        sb.append(content);

        //return sb.toString();
        return connector.askWithHistory(sb.toString());
    }

    public String doSplitNoWithHistory(String content) {
        StringBuilder sb = new StringBuilder();
        sb.append(prompt);
        sb.append(content);
        List<ChatMessage> messages = new ArrayList<>();
        messages.add(new ChatMessage(ChatMessageRole.USER.value(), sb.toString()));
        return connector.ask(messages.toString());
    }

    public static String create(String content) {
        boolean head = true;
        int num = 0;
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < content.length(); i++) {
            //System.out.println(content.charAt(i)+" "+(content.charAt(i)=='“'?"=":"!="));
            if (content.charAt(i) == '”' || content.charAt(i) == '“' || content.charAt(i) == '"') {
                if (head) {
                    sb.append('\n');
                    sb.append("[" + num + "]");
                    sb.append(content.charAt(i));
                    head = false;
                    num++;
                } else {
                    sb.append(content.charAt(i));
                    sb.append('\n');
                    head = true;
                }
            } else {
                sb.append(content.charAt(i));
            }
        }
        return sb.toString();
    }

    /**
     */
    static public Map<String, String> getVoiceRoles(String roles_string,Map<String, String> role_voice_map){
        String[] roles = roles_string.split("\n");
        if(role_voice_map==null) role_voice_map = new HashMap<>();
        for (int i = 0; i < roles.length; i++) {
            String s = roles[i];
            String[] split = s.split("[：:]");
            int length = split.length;
            Integer key = i;
            if (StringUtils.isNumeric(split[0])) key = Integer.parseInt(split[0]);
            role_voice_map.put(split[Math.min(length-1, 1)], split[Math.min(length-1, 2)]);
        }
        role_voice_map.put("旁白", "男性主持人");
        return role_voice_map;
    }
    public static List<Dialog> refacterDialogs(String content,String keyword,Map<String, String> role_voice_map){
        String[] keywords=keyword.split("[：\n]");
        List<Dialog> ls=new ArrayList<>();
        boolean head=true;
        int num=0;
        int begin=0;
        for(int i=0; i<content.length();i++){
            if(content.charAt(i)=='”'||content.charAt(i)=='“'||content.charAt(i)=='"'){
                if(head){
                    if(begin<i) {
                        ls.add(new Dialog(content.substring(begin, i), "旁白", role_voice_map.get("旁白")));
                    }
                    begin=i;
                    head=false;}
                else{
                    ls.add(new Dialog(content.substring(begin, i+1),keywords[num*3+1], role_voice_map.get(keywords[num*3+1])));
                    begin=i+1;
                    head=true;num++;}
            }
        }
        if(begin<content.length())ls.add(new Dialog(content.substring(begin, content.length()),"旁白", role_voice_map.get("旁白")));
        return ls;
    }

    public static List<Dialog> splitDialogs(String content,String keyword,String role_voice){
        String[] keywords=keyword.split("[：\n]");
        List<Dialog> ls=new ArrayList<>();
        boolean head=true;
        int num=0;
        int begin=0;
        for(int i=0; i<content.length();i++){
            if(content.charAt(i)=='”'||content.charAt(i)=='“'||content.charAt(i)=='"'){
                if(head){
                    if(begin<i) {
                        ls.add(new Dialog(content.substring(begin, i), "旁白", role_voice));
                    }
                    begin=i;
                    head=false;}
                else{
                    ls.add(new Dialog(content.substring(begin, i+1),keywords[num*3+1], role_voice));
                    begin=i+1;
                    head=true;num++;}
            }
        }
        if(begin<content.length())ls.add(new Dialog(content.substring(begin, content.length()),"旁白", role_voice));
        return ls;
    }
    static public List<Dialog> getDialogs(String text, String roles_string,Map<String, String> role_voice_map) {
        String[] roles = roles_string.split("\n");
        Map<Integer, String> role_map = new HashMap<>();

        for (int i = 0; i < roles.length; i++) {
            String s = roles[i];
            String[] split = s.split("[：:]");
            int length = split.length;
            Integer key = i;
            if (StringUtils.isNumeric(split[0])) key = Integer.parseInt(split[0]);
            role_map.put(key, split[Math.min(length-1, 1)]);
            role_voice_map.put(split[Math.min(length-1, 1)], split[Math.min(length-1, 2)]);
        }

        role_voice_map.put("旁白", "男性主持人");

        System.out.println(role_map);
        System.out.println(role_voice_map);

        LinkedList<Dialog> dialogs = new LinkedList<>();

        String[] targetText = text.split("\n");
        Integer key = 0;
        for (int i = 0; i < targetText.length; i++) {
            String s = targetText[i];
            Dialog dialog = new Dialog();
            if (StringUtils.isBlank(s)) continue;
            int left = s.indexOf('['), right = s.indexOf(']');
            if (left == -1) {
                dialog.setRole("旁白");
                dialog.setText(s);
                String vg = role_voice_map.get("旁白");
                if (vg == null) {
                    vg = role_voice_map.get("\"旁白\"");
                }
                dialog.setVoice_generator(vg);
                dialogs.add(dialog);
                continue;
            }
            String substring = s.substring(left + 1, right - 1);
            boolean ifNumeric = StringUtils.isNumeric(substring);

            if (ifNumeric) key = Integer.parseInt(substring);

            System.out.println(key);
            String role = role_map.get(key);

            String vg = role_voice_map.get(role);
            if (vg == null) {
                vg = role_voice_map.get("\"" + role + "\"");
            }

            dialog.setRole(role);
            dialog.setText(s.substring(right+1));
            dialog.setVoice_generator(vg);

            dialogs.add(dialog);
            key++;
        }

        return dialogs;
    }

    public static void main(String[] args) {
        LLMSplitTool tool = new LLMSplitTool(
                new BlueLM()
        );

        String targetText = LLMSplitTool.create("无视于身后那阴森森的目光，萧炎抱着卷轴在斗技堂管理员处登记了一下，这才与薰儿轻声笑谈着，慢吞吞的行出了斗技堂。\n" +
                "\n" +
                "“小混蛋，你给我等着吧，等你被分配出家族之后，我有的是时间收拾你！失去了族长的庇护，你狗屁都不是！”望着那逐渐远去的背影，萧宁恨得有些牙痒痒，咬牙切齿的反手一掌轰在身旁的书架上，顿时在其上留下了一个浅浅的手印。\n" +
                "\n" +
                "行出斗技堂，萧炎先是陪着兴致勃勃的薰儿到后山逛了一下午，待到天色渐暗之后，这才回到自己的小窝。\n" +
                "\n" +
                "回到房间，关上房门，萧炎双肩顿时跨了下来，将卷轴放在桌上，端起茶杯，一饮而尽，有些后怕的苦笑道：“这妮子，真是太能走了。”\n" +
                "\n" +
                "“那小丫头，来历似乎有点不一般啊。”苍老的声音，忽然的在房间中响起。\n" +
                "\n" +
                "有气无力的抬了抬眼，望着那犹如鬼魅一般出现在房间中的药老，萧炎撇了撇嘴，懒懒的问道：“老师知道她的来历？”");
        Map<String, String> map = new HashMap<String, String>();
        map.put("萧炎","精英男声");
        map.put("药老","男性主持人A");
        String s = tool.doSplitWithHistory(targetText,map);
        System.out.println("--------------------------");
        System.out.println(s);
        System.out.println("--------------------------");
        map=getVoiceRoles(s,map);
        List<Dialog> dialogs = tool.getDialogs(targetText, s,map);
        System.out.println(dialogs);
    }
}
